import os
import re
import shutil
from collections import defaultdict
from pathlib import Path

def check_and_delete_missing_modalities(root_dir, dry_run=True, log_file=None):
    """
    检查指定目录下的图片是否存在缺失的模态，并可选择删除不完整的样本
    
    参数:
    root_dir (str): 根目录路径，包含多个类别子文件夹
    dry_run (bool): 是否执行干运行（只显示操作不实际删除）
    log_file (str): 日志文件路径，用于记录操作
    """
    # 用于存储每个样本的两种模态图片
    samples = defaultdict(lambda: {"TrueLight": None, "SolderLight": None})
    
    # 正则表达式模式，用于提取样本名称和模态类型
    pattern = re.compile(r'^(.*?)(TrueLight|SolderLight)(\.[a-zA-Z0-9]+)$')
    
    # 遍历所有子文件夹和文件
    for class_dir in os.listdir(root_dir):
        class_path = os.path.join(root_dir, class_dir)
        if not os.path.isdir(class_path):
            continue
        
        for filename in os.listdir(class_path):
            file_path = os.path.join(class_path, filename)
            if not os.path.isfile(file_path):
                continue
            
            # 尝试匹配文件名模式
            match = pattern.match(filename)
            if match:
                sample_name, modality, extension = match.groups()
                # 存储样本的模态信息
                samples[sample_name][modality] = file_path
    
    # 找出仅存在一种模态的样本
    missing_true_light = []
    missing_solder_light = []
    
    for sample_name, modalities in samples.items():
        if modalities["TrueLight"] and not modalities["SolderLight"]:
            missing_solder_light.append(modalities["TrueLight"])
        elif modalities["SolderLight"] and not modalities["TrueLight"]:
            missing_true_light.append(modalities["SolderLight"])
    
    # 准备日志记录
    log_messages = []
    
    # 删除不完整的样本
    deleted_count = 0
    for file_path in missing_true_light + missing_solder_light:
        action = "将删除" if dry_run else "已删除"
        log_msg = f"{action}: {file_path}"
        log_messages.append(log_msg)
        
        if not dry_run:
            try:
                os.remove(file_path)
                deleted_count += 1
            except Exception as e:
                log_messages.append(f"删除失败: {file_path} - {str(e)}")
    
    # 输出统计结果
    log_messages.append("\n=== 统计结果 ===")
    log_messages.append(f"仅存在SolderLight的样本: {len(missing_true_light)}")
    log_messages.append(f"仅存在TrueLight的样本: {len(missing_solder_light)}")
    log_messages.append(f"总共删除的文件数: {deleted_count if not dry_run else '未执行删除（干运行模式）'}")
    
    # 打印日志
    for msg in log_messages:
        print(msg)
    
    # 如果指定了日志文件，则写入日志
    if log_file:
        with open(log_file, 'w') as f:
            f.write('\n'.join(log_messages))
        print(f"\n日志已保存至: {log_file}")
    
    return len(missing_true_light), len(missing_solder_light), deleted_count

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='检查并删除仅存在单一模态的图片')
    parser.add_argument('--dataset', default='data_zoo/PCBA/20241209_PCBA样板标图数据含真彩-华科用',help='数据集根目录路径')
    parser.add_argument('--delete', default='True',action='store_true', help='执行实际删除操作（默认仅显示不删除）')
    parser.add_argument('--log', default='modalcheck.log',help='日志文件路径')
    
    args = parser.parse_args()
    
    print("=== 开始检查不完整的模态样本 ===")
    if not args.delete:
        print("警告: 当前处于干运行模式，不会执行实际删除操作。使用 --delete 参数执行删除。")
    
    # 检查并删除缺失的模态
    missing_true, missing_solder, deleted = check_and_delete_missing_modalities(
        args.dataset, 
        dry_run=not args.delete,
        log_file=args.log
    )
    
    print("=== 操作完成 ===")    

# if __name__ == "__main__":
#     import argparse
    
#     # parser = argparse.ArgumentParser(description='检查图片是否存在缺失的模态')
#     # parser.add_argument('--dataset',  help='数据集根目录路径')
#     # parser.add_argument('--output', help='输出结果文件路径')
    
#     # args = parser.parse_args()
#     # args.dataset='"F:\\jrs项目\\20241209_PCBA样板标图数据-华科用\\20241209_PCBA样板标图数据含真彩-华科用'
    
#     # 检查缺失的模态
#     missing_true_light, missing_solder_light = check_missing_modalities('F:\\jrs项目\\20241209_PCBA样板标图数据-华科用\\20241209_PCBA样板标图数据含真彩-华科用')
    
#     # 输出结果
#     print("\n=== 仅存在SolderLight模态的图片 ===")
#     for filename in missing_true_light:
#         print(filename)
    
#     print("\n=== 仅存在TrueLight模态的图片 ===")
#     for filename in missing_solder_light:
#         print(filename)
    
#     print(f"\n统计结果:")
#     print(f"仅存在SolderLight的样本: {len(missing_true_light)}")
#     print(f"仅存在TrueLight的样本: {len(missing_solder_light)}")
    
#     # 如果指定了输出文件，则写入结果
#     output_path='modalcheck.log'
#     with open(output_path, 'w') as f:
#         f.write("=== 仅存在SolderLight模态的图片 ===\n")
#         for filename in missing_true_light:
#             f.write(f"{filename}\n")
        
#         f.write("\n=== 仅存在TrueLight模态的图片 ===\n")
#         for filename in missing_solder_light:
#             f.write(f"{filename}\n")
        
#         f.write(f"\n统计结果:\n")
#         f.write(f"仅存在SolderLight的样本: {len(missing_true_light)}\n")
#         f.write(f"仅存在TrueLight的样本: {len(missing_solder_light)}\n")
    
#     print(f"\n结果已保存至: {output_path}")    